#include <gtest/gtest.h>
#include "cipher.hpp"
#include <filesystem>

class CipherTest : public ::testing::Test {
  protected:
    std::string temp_dir;
    std::string prvfile;
    std::string pucfile;

    void SetUp() override
    {
        temp_dir = "/tmp/syralink_test_key";
        prvfile  = temp_dir + "/private.pem";
        pucfile  = temp_dir + "/public.pem";

        if(std::filesystem::exists(temp_dir) == true) {
            return;
        }
        std::filesystem::create_directory(temp_dir);

        std::string prv_cmd =
            "openssl genpkey -algorithm EC -pkeyopt ec_paramgen_curve:SM2 -out " +
            prvfile;
        std::string puc_cmd =
            "openssl pkey -in " + prvfile + " -pubout -out " + pucfile;
        run_cmd(prv_cmd);
        run_cmd(puc_cmd);
    }

    void TearDown() override { }

    int run_cmd(const std::string& cmd)
    {
        return system(cmd.c_str());
    }
};

TEST_F(CipherTest, SM3_CorrectTest)
{
    std::vector<sl::byte> sm3_std = sl::str2bin(
        "66c7f0f462eeedd9d1f2d46bdc10e4e24167c4875cf2f7a2297da02b8f4ba8e0");
    for(size_t i = 0; i < 1000; i++) {
        std::vector<sl::byte> hash = sl::calculate_hash({'a', 'b', 'c'});
        ASSERT_EQ(hash.size(), 32);
        ASSERT_EQ(hash, sm3_std);
    }
}

TEST_F(CipherTest, SM2_ReadKeyFile)
{
    for(size_t i = 0; i < 1000; i++) {
        std::vector<sl::byte> prv = sl::read_private(prvfile);
        std::vector<sl::byte> puc = sl::read_public(pucfile);

        ASSERT_EQ(prv.size(), 32);
        ASSERT_EQ(puc.size(), 64);

        std::vector<sl::byte> data(66, 1);
        std::vector<sl::byte> sign = sl::ecc_sign(data, prv);
        ASSERT_EQ(sign.size(), 64);
        ASSERT_TRUE(sl::ecc_verify(data, sign, puc));
    }
}

TEST_F(CipherTest, SM2_SignVerify)
{
    std::vector<sl::byte> prv(128);
    std::vector<sl::byte> puc(128);

    for(size_t i = 0; i < 1000; i++) {
        bool rc = sl::ecc_generate(prv, puc);
        ASSERT_TRUE(rc);

        std::vector<sl::byte> data(66, 1);
        std::vector<sl::byte> sign = sl::ecc_sign(data, prv);
        ASSERT_EQ(sign.size(), 64);

        rc = sl::ecc_verify(data, sign, puc);
        ASSERT_TRUE(rc);
    }
}

TEST_F(CipherTest, SM2_Encrypt)
{
    std::vector<sl::byte> prv(128);
    std::vector<sl::byte> puc(128);

    for(size_t i = 0; i < 1000; i++) {
        bool rc = sl::ecc_generate(prv, puc);
        ASSERT_TRUE(rc);

        std::vector<sl::byte> data(32, 1);
        std::vector<sl::byte> enc = sl::ecc_encrypt(data, puc);
        ASSERT_EQ(enc.size(), 64 + 32 + data.size());

        std::vector<sl::byte> dec = sl::ecc_decrypt(enc, prv);
        ASSERT_EQ(dec, data);
    }
}

TEST_F(CipherTest, SM4_Encrypt)
{
    std::vector<sl::byte> data(32, 1);
    std::vector<sl::byte> key(16, 2);

    for(size_t i = 0; i < 1500; i++) {
        std::vector<sl::byte> enc = sl::sym_encrypt(data, key);
        std::vector<sl::byte> dec = sl::sym_decrypt(enc, key);
        ASSERT_EQ(dec, data);
    }

    std::vector<sl::byte> data2(9, 2);
    for(size_t i = 0; i < 1500; i++) {
        std::vector<sl::byte> enc = sl::sym_encrypt(data2, key);
        std::vector<sl::byte> dec = sl::sym_decrypt(enc, key);
        ASSERT_EQ(dec, data2);
    }

    std::vector<sl::byte> data3(21, 3);
    for(size_t i = 0; i < 1500; i++) {
        std::vector<sl::byte> enc = sl::sym_encrypt(data3, key);
        std::vector<sl::byte> dec = sl::sym_decrypt(enc, key);
        ASSERT_EQ(dec, data3);
    }
}
